package sock

import (
	"io"
	"net"
	"time"
)

// 带有超时机制的连接
// 可以设置读取超时时间，单位为毫秒
// 可以设置编码、解码函数，用于加密、解密数据
// 可以设置连接超时时间，单位为毫秒
type TimeoutConn struct {
	net.Conn
	timeout time.Duration
	peek    []byte
	readBuf []byte              // 读取缓冲区，不可并发读取
	encode  func([]byte) []byte // 编码函数, 用于加密数据
	decode  func([]byte) []byte // 解码函数, 用于解密数据
}

// 带有超时机制的连接
// network: 网络类型, 如: tcp, udp
// timeouts: 超时时间, 单位为毫秒,第一个为读取超时时间，默认60秒; 第二个为连接超时时间, 默认为10秒
// 返回值: 连接对象, 错误信息
func DialTimeoutConn(network, address string, timeoutMss ...int64) (conn *TimeoutConn, err error) {
	timeoutMs := int64(60000)     // 读取超时时间, 单位为毫秒, 默认为60000毫秒
	dialTimeoutMs := int64(10000) // 连接超时时间, 单位为毫秒, 默认为10000毫秒
	if len(timeoutMss) > 0 {
		timeoutMs = timeoutMss[0]
	}
	if len(timeoutMss) > 1 {
		dialTimeoutMs = timeoutMss[1]
	}
	timeoutMs = min(max(timeoutMs, 1000), 300000)
	dialTimeoutMs = min(max(dialTimeoutMs, 50), 60000)
	c1, err := net.DialTimeout(network, address, time.Millisecond*time.Duration(dialTimeoutMs))
	if err != nil {
		return nil, err
	}

	if c2, ok := c1.(*net.TCPConn); ok {
		//设置tcp心跳时间
		c2.SetKeepAlive(true)
		c2.SetKeepAlivePeriod(60 * time.Second)
	}
	conn = NewTimeoutConn(c1, timeoutMs)
	return conn, nil
}

//  DialTimeoutTCP  dial tcp 超时时间, 单位为毫秒, 默认为30000毫秒
func DialTimeoutTCP( address string, timeoutMss ...int64) (conn *TimeoutConn, err error){
	return DialTimeoutConn("tcp", address, timeoutMss...)
}

// 创建带有超时机制的连接
// timeoutMs: 超时时间, 单位为毫秒
func NewTimeoutConn(conn net.Conn, timeoutMs int64) *TimeoutConn {
	tr := new(TimeoutConn)
	tr.Conn = conn
	tr.timeout = time.Millisecond * time.Duration(timeoutMs)
	return tr
}

// 设置加密、解密函数
// encode: 加密函数, 用于加密数据
// decode: 解密函数, 用于解密数据
// 返回值: 连接对象
func (m *TimeoutConn) SetCryption(encode, decode func([]byte) []byte) *TimeoutConn {
	m.encode = encode
	m.decode = decode
	return m
}

// 带有超时机制的读取函数
// data: 读取的数据
func (m *TimeoutConn) Read(data []byte) (num int, err error) {
	if len(m.peek) > 0 {
		num = copy(data, m.peek)
		m.peek = m.peek[num:]
		return
	}
	err = m.Conn.SetReadDeadline(time.Now().Add(m.timeout))
	if err != nil {
		return
	}
	num, err = m.Conn.Read(data)
	return
}

// 带有超时机制的读取函数, 读取指定的字节数, 当没有错误发生时, num = len(data)
// data: 读取的数据
func (m *TimeoutConn) ReadFull(data []byte) (num int, err error) {
	return io.ReadFull(m, data)
}

// 带有超时机制的写入函数
// data: 写入的数据
func (m *TimeoutConn) Write(data []byte) (num int, err error) {
	num, err = m.Conn.Write(data)
	return
}

// 关闭连接
func (m *TimeoutConn) Close() (err error) {
	err = m.Conn.Close()
	return
}

// Peek 查看指定长度的字节
// n: 查看的字节数
func (m *TimeoutConn) Peek(n int) (data []byte, err error) {
	if len(m.peek) > 0 {
		data = m.peek
		return
	}

	data = make([]byte, n)
	num, err := m.Read(data)
	if err != nil {
		return nil, err
	}
	m.peek = data[:num]
	return m.peek, nil
}

// 检查读取缓冲区长度, 并调整缓冲区长度, 确保缓冲区长度大于等于l
func (m *TimeoutConn) checkReadBufLen(l int) {
	if cap(m.readBuf) < l {
		m.readBuf = make([]byte, l)
	}
	if len(m.readBuf) < l {
		m.readBuf = m.readBuf[:l]
	}
}

// ReadLenString 读取指定长度的字符串
/*
* 格式: 2字节长度 + 数据
 */
func (m *TimeoutConn) ReadLenBytes() (data []byte, err error) {
	m.checkReadBufLen(2)
	n, err := m.ReadFull(m.readBuf[:2])
	if err != nil {
		return
	}
	if n != 2 {
		err = io.ErrUnexpectedEOF
		return
	}

	l := uint16(m.readBuf[0])<<8 | uint16(m.readBuf[1])
	m.checkReadBufLen(int(l))
	n, err = m.ReadFull(m.readBuf[:l])
	if n == 0 {
		return
	}
	if m.decode != nil {
		data = m.decode(m.readBuf[:n])
	} else {
		data = m.readBuf[:n]
	}
	return
}

// WriteLenBytes 写入指定长度的字节
/*
* 格式: 2字节长度 + 数据
 */
func (m *TimeoutConn) WriteLenBytes(data []byte) (n int, err error) {
	if m.encode != nil {
		data = m.encode(data)
	}
	b := make([]byte, 2)
	l := len(data)
	b[0] = byte(uint16(l) >> 8)
	b[1] = byte(l)
	_, err = m.Conn.Write(b)
	if err != nil {
		n = 0
		return
	}
	n, err = m.Conn.Write(data)
	return
}
